# Dickstein, Ho, and Mark (2023)
# This script runs the counterfactuals for each of the bootstrap draws.
# First, I create the matrix of parameters based on the bootstrap draws
# Second, I run the counterfactuals using this matrix

# * # * # * # * # * # * #
# PRELIMINARIES         #
# * # * # * # * # * # * #
setwd("../../../library")
source("PreliminariesCode.R")

find_names <- function(name_vec, start){
  unlist(lapply(strsplit(name_vec[which(type_param_vec %in% start)], "\\."), function(x) x[3]))
}

# What bootstrap draws in param_mat 
k_vec <- 1:20

# What bootstrap draws to simulate outcomes
k_vec_sim <- 1:20

# Are you providing your own parameters? (if yes, set to F)
load_params = F

# Save the counterfactual results as you go?
save_asu_go <- T

# Which rating areas?
ravec = 1:7

# Which years?
yrvec = 2016

# What percent of the premium is covered by an employer?
nu = .65

# What is the behavioral share
beshr <- .01

# What alpha-psi cutoff to use: 
alphampsicutoff <- .05

# What ACG cutoff to use: 
acgHL_boundary <- 1.76

# What is the base markup?
mrkup_0 <- 0

# What is psi_upperbound?
psi_upperbound <- .3

# Which spec do you want to use?
spectouse <- "lowrisk_simplemh_censor0.2"

# Where is the counterfactual folder?
counterfactual_folder <- paste0(project_folder, "/analysis/counterfactuals")

# Where is the output folder
output_folder <- paste0(project_folder, "/analysis/counterfactuals/bootstrap/output")
dir.create(output_folder, recursive = T)

# where is datasets folder
data_folder <- paste0(project_folder, "/data")

# Loading counterfactual functions 
source(paste0(counterfactual_folder, "/library/DA03aRunCounterfactualFunctions.R"))

# Which counterfactuals do you want to run? 
run_which_ctfls <- c(0, 1, 3, 5)

# Which columns do you want to create? 
create_which_columns <- c(0, 1, 3, 5)

# How many parameters are there? 
n_pars <- 41

# ~ # ~ # V CREATING PARAMETER MATRIX V # ~ # ~ #
param_mat <- matrix(0, nrow = length(k_vec), ncol = n_pars)
for(k in k_vec){

  # ~ # ~ # ~ # ~ # ~ #
  # Upload parameters #
  # ~ # ~ # ~ # ~ # ~ #
   
  # Loading individual market demand params
  IndMarketEstimates_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/lowrisk_simplemh_censor0.2_main/bootstrap/output_", k, "/solution.csv"))
  
  ## Kondo.
  IndMarketEstimates_0 <- IndMarketEstimates_0[, -c(1,2)]
  
  # Loading the number of beta1, 2, 3, 0 parameters:
  IndMarketlabels_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/lowrisk_simplemh_censor0.2_main/bootstrap/output_", k, "/covar_lens.csv"))
  IndMarSE_0 <- cumsum(IndMarketlabels_0$x)
  
  # Putting them into an easy to read data table
  NamedBetaList_0 <- list(
    beta1 = as.numeric(IndMarketEstimates_0[1, 1:IndMarSE_0[1]]),
    beta2 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[1]):IndMarSE_0[2]]),
    beta3 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[2]):IndMarSE_0[3]]),
    beta0 = as.numeric(IndMarketEstimates_0[1, (1 + IndMarSE_0[3]):IndMarSE_0[4]])
  )
  # Adding the original variable names back to reorganized nicely beta list
  names(NamedBetaList_0[[1]]) <- names(IndMarketEstimates_0)[1:IndMarSE_0[1]]
  names(NamedBetaList_0[[2]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[1]):IndMarSE_0[2]]
  names(NamedBetaList_0[[3]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[2]):IndMarSE_0[3]]
  names(NamedBetaList_0[[4]]) <- names(IndMarketEstimates_0)[(1 + IndMarSE_0[3]):IndMarSE_0[4]]
  
  # Do the same thing for the switcher beta variables
  SGMarketEstimates_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/lowrisk_simplemh_censor0.2_switcher/bootstrap/output_", k, "/solution.csv"))
  SGMarketEstimates_0 <- SGMarketEstimates_0[, -c(1,2)]
  SGMarketlabels_0 <- fread(paste0(project_folder, "/analysis/estimationcode/specs/lowrisk_simplemh_censor0.2_switcher/bootstrap/output_", k, "/covar_lens.csv"))
  SGMarSE_0 <- cumsum(SGMarketlabels_0$x)
  
  # Putting them into an easy to read data table
  SGNamedBetaList_0 <- list(
    beta1 = as.numeric(SGMarketEstimates_0[1, 1:SGMarSE_0[1]]),
    beta2 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[1]):SGMarSE_0[2]]),
    beta3 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[2]):SGMarSE_0[3]]),
    beta0 = as.numeric(SGMarketEstimates_0[1, (1 + SGMarSE_0[3]):SGMarSE_0[4]])
  )
  names(SGNamedBetaList_0[[1]]) <- names(SGMarketEstimates_0)[1:SGMarSE_0[1]]
  names(SGNamedBetaList_0[[2]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[1]):SGMarSE_0[2]]
  names(SGNamedBetaList_0[[3]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[2]):SGMarSE_0[3]]
  names(SGNamedBetaList_0[[4]]) <- names(SGMarketEstimates_0)[(1 + SGMarSE_0[3]):SGMarSE_0[4]]
  
  # Loading supply size parameters 
  SupplySideCoefficients_0 <-fread(paste0("~/sharedWork/oregon/Analysis_NDM/prem_setting/bootstrap/release/pr_paperfull_best_guess_ra_s_wsimcost_", k, ".csv"))
  SupplySideData_0 <- data.table(
    expand.grid(
      payer_id = c("M0013", "M0019", "M0035", "M0062", "M0077", "M0080"),
      year = c("2015", "2016"))
  )
  
  # Create beta_{5j}'A_j
  names(SupplySideCoefficients_0) <- c("var", "mod1", "mod2", "mod3", "mod4") 
  SupplySideCoefficients_0[, year := ifelse(grepl("2015", var), "2015", NA), ]
  SupplySideCoefficients_0[, year := ifelse(grepl("2016", var), "2016", year), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0013", var), "M0013", NA), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0019", var), "M0019", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0035", var), "M0035", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0062", var), "M0062", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0077", var), "M0077", payer_id), ]
  SupplySideCoefficients_0[, payer_id := ifelse(grepl("M0080", var), "M0080", payer_id), ]
  constant_0 <- as.numeric(SupplySideCoefficients_0[var == "Constant"]$mod3)
  
  SupplySideCoefficients2015_0 <- SupplySideCoefficients_0[year == "2015" | is.na(year)]
  SupplySideCoefficients2015_0 <- SupplySideCoefficients2015_0[ , .(
    fe_beta_5 = sum(as.numeric(mod3)) + constant_0,
    year = "2015"
  ), by = c("payer_id")]
  SupplySideCoefficients2016_0 <- SupplySideCoefficients_0[year == "2016" | is.na(year)]
  SupplySideCoefficients2016_0 <- SupplySideCoefficients2016_0[ , .(
    fe_beta_5 = sum(as.numeric(mod3)) + constant_0,
    year = "2016"
  ), by = c("payer_id")]
  SupplySideFEs_0 <- rbind(SupplySideCoefficients2015_0, SupplySideCoefficients2016_0)
  SupplySideData_0 <- merge(
    SupplySideData_0,
    SupplySideFEs_0,
    by = c("payer_id", "year"),
    all.x = T
  )
  
  # Add beta_4
  SupplySideData_0$beta_4 <- as.numeric(SupplySideCoefficients_0[var == "Medical costs"]$mod3)
  
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  # Create parameter value matrix #
  # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
  param_mat[k, ] <- c(
    unlist(NamedBetaList_0), 
    unlist(SGNamedBetaList_0), 
    SupplySideData_0$fe_beta_5, 
    SupplySideData_0$beta_4[1],
    mrkup_0)
  
  if(k == k_vec[1]){
    param_mat_names <- c(
      paste0("im.", names(unlist(NamedBetaList_0))),
      paste0("sg.", names(unlist(SGNamedBetaList_0))),
      paste0("al.beta5.", SupplySideData_0$payer_id, "_", SupplySideData_0$year),
      "al.beta4.constant",
      "mrkup")
  } else {
    if(all(param_mat_names != c(
      paste0("im.", names(unlist(NamedBetaList_0))),
      paste0("sg.", names(unlist(SGNamedBetaList_0))),
      paste0("al.beta5.", SupplySideData_0$payer_id, "_", SupplySideData_0$year),
      "al.beta4.constant",
      "mrkup"))){ print("Warning: variable names do not match for different bootstrap draws.")} 
  }
}

param_mat <- as.data.table(param_mat)
names(param_mat) <- param_mat_names

# ~ # ~ # ^ IF YOU ARE CREATING PARAM MAT HERE ^ # ~ # ~ #

# ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #
# For each row of the parameter value matrix, run a counterfactual  #
# ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ # ~ #

results_list_base <- list()
results_list_1 <- list()
results_list_3 <- list()
results_list_5 <- list()

for(i in k_vec_sim){
  print(paste0("ESTIMATING COUNTERFACTUAL", i))
  
  # Unpack the matrix into NamedBetaList SGNamedBetaList SupplySidedata and mrkup 
  param_vec <- as.vector(t(param_mat[i, ]))
  names(param_vec) <- names(param_mat)
  type_param_vec <- substr(names(param_mat), 1, 8)
  
  # NamedBetaList
  NamedBetaList <- list(
    beta1 = param_vec[which(type_param_vec %in% "im.beta1")],
    beta2 = param_vec[which(type_param_vec %in% "im.beta2")],
    beta3 = param_vec[which(type_param_vec %in% "im.beta3")],
    beta0 = param_vec[which(type_param_vec %in% "im.beta0")]
  )
  names(NamedBetaList[[1]]) <- find_names(names(param_vec), "im.beta1")
  names(NamedBetaList[[2]]) <- find_names(names(param_vec), "im.beta2")
  names(NamedBetaList[[3]]) <- find_names(names(param_vec), "im.beta3")
  names(NamedBetaList[[4]]) <- find_names(names(param_vec), "im.beta0")
  
  #SGNamedBetaList
  SGNamedBetaList <- list(
    beta1 = param_vec[which(type_param_vec %in% "sg.beta1")],
    beta2 = param_vec[which(type_param_vec %in% "sg.beta2")],
    beta3 = param_vec[which(type_param_vec %in% "sg.beta3")],
    beta0 = param_vec[which(type_param_vec %in% "sg.beta0")]
  )
  names(SGNamedBetaList[[1]]) <- find_names(names(param_vec), "sg.beta1") 
  names(SGNamedBetaList[[2]]) <- find_names(names(param_vec), "sg.beta2")
  names(SGNamedBetaList[[3]]) <- find_names(names(param_vec), "sg.beta3")
  names(SGNamedBetaList[[4]]) <- find_names(names(param_vec), "sg.beta0")
  
  # SupplySidedata
  SupplySideData <- data.table(
    payer_id = substr(find_names(names(param_vec), "al.beta5"), 1, 5), 
    year = str_sub(find_names(names(param_vec), "al.beta5"), -4, -1),
    fe_beta_5 = param_vec[which(type_param_vec %in% "al.beta5")], 
    beta_4 = param_vec[which(type_param_vec %in% "al.beta4")]
  )

  # Markup
  mrkup <- param_vec["mrkup"]
  
  # Run the data readying file
  source(paste0(counterfactual_folder, "/DA03bRunCounterfactualGetdataready.R"))
  
  # Run the relevant counterfactuals
  tryCatch(
    source(paste0(counterfactual_folder, "/specs/", spectouse, "/DA03dCounterfactualCode.R")),
      error = function(e){
        resbase <- NA
        res1 <- NA
        res3 <- NA
        res5 <- NA
      }
    )

  # Save as you go
  if(save_asu_go == T){
    save(resbase, file = paste0(output_folder, "/output_base_results_draw", i, ".Rdata"))
    save(res1, file = paste0(output_folder, "/output_1_results_draw",i, ".Rdata"))
    save(res3, file = paste0(output_folder, "/output_3_results_draw", i, ".Rdata"))
    save(res5, file = paste0(output_folder, "/output_5_results_draw", i, ".Rdata")) 
  }
  
  #Kondo!
  rm(ExplodedData_SG) 
  rm(ExplodedData_SG_small)
  rm(ExplodedData_Ind)
  rm(ExplodedData_Ind_small)
}
